package middleware

import (
	"context"
	"fmt"
	"strconv"
	"time"

	"gitee.com/war2001/NiuMaEasyGo/global"
	"github.com/gin-gonic/gin"
	"github.com/juju/ratelimit"
	"go.uber.org/zap"
)

// RateLimiterType 限流类型
type RateLimiterType string

const (
	// GlobalLimiter 全局限流器
	GlobalLimiter RateLimiterType = "global"
	// IPLimiter 基于IP的限流器
	IPLimiter RateLimiterType = "ip"
	// UserLimiter 基于用户的限流器
	UserLimiter RateLimiterType = "user"
	// APILimiter 基于API的限流器
	APILimiter RateLimiterType = "api"
)

// RateLimiterConfig 限流配置
type RateLimiterConfig struct {
	// 限流类型
	Type RateLimiterType
	// 填充速率(每秒填充的令牌数)
	FillRate float64
	// 桶容量
	Capacity int64
	// 键前缀
	KeyPrefix string
	// 限流时间窗口(秒)
	Window int
	// 是否使用Redis
	UseRedis bool
}

// RateLimiter 限流中间件
func RateLimiter(config RateLimiterConfig) gin.HandlerFunc {
	// 如果使用Redis，则使用Redis实现的限流器
	if config.UseRedis && global.Redis != nil {
		return redisRateLimiter(config)
	}

	// 否则使用本地令牌桶实现的限流器
	return localRateLimiter(config)
}

// localRateLimiter 本地令牌桶限流器
func localRateLimiter(config RateLimiterConfig) gin.HandlerFunc {
	// 创建一个令牌桶
	bucket := ratelimit.NewBucketWithRate(config.FillRate, config.Capacity)

	return func(c *gin.Context) {
		// 根据限流类型获取限流键
		// 注意：本地令牌桶实现中，我们不需要使用key，因为使用的是单一的全局桶
		// 这里只是为了保持代码一致性，实际上没有使用这个key
		_ = getLimiterKey(c, config)

		// 尝试获取令牌
		if bucket.TakeAvailable(1) < 1 {
			// 没有可用的令牌，返回429错误
			ResponseTooManyRequests(c, "请求过于频繁，请稍后再试")
			c.Abort()
			return
		}

		c.Next()
	}
}

// redisRateLimiter Redis实现的限流器
func redisRateLimiter(config RateLimiterConfig) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 根据限流类型获取限流键
		key := getLimiterKey(c, config)
		limiterKey := fmt.Sprintf("%s:%s", config.KeyPrefix, key)

		// 创建上下文
		ctx := context.Background()

		// 获取当前计数
		countStr, err := global.Redis.Get(ctx, limiterKey).Result()
		if err != nil && err.Error() != "redis: nil" {
			zap.L().Error("获取限流计数失败", zap.Error(err))
			c.Next()
			return
		}

		// 如果键不存在，则设置初始值为1，并设置过期时间
		if countStr == "" {
			err = global.Redis.Set(ctx, limiterKey, 1, time.Duration(config.Window)*time.Second).Err()
			if err != nil {
				zap.L().Error("设置限流计数失败", zap.Error(err))
			}
			c.Next()
			return
		}

		// 将计数转换为整数
		count, err := strconv.Atoi(countStr)
		if err != nil {
			zap.L().Error("转换限流计数失败", zap.Error(err))
			c.Next()
			return
		}

		// 检查是否超过限制
		if count >= int(config.Capacity) {
			// 超过限制，返回429错误
			ResponseTooManyRequests(c, "请求过于频繁，请稍后再试")
			c.Abort()
			return
		}

		// 增加计数
		_, err = global.Redis.Incr(ctx, limiterKey).Result()
		if err != nil {
			zap.L().Error("增加限流计数失败", zap.Error(err))
		}

		c.Next()
	}
}

// getLimiterKey 获取限流键
func getLimiterKey(c *gin.Context, config RateLimiterConfig) string {
	switch config.Type {
	case IPLimiter:
		return c.ClientIP()
	case UserLimiter:
		// 从上下文中获取用户ID
		userID, exists := c.Get("user_id")
		if !exists {
			return "anonymous"
		}
		return fmt.Sprintf("user:%v", userID)
	case APILimiter:
		return fmt.Sprintf("api:%s:%s", c.Request.Method, c.FullPath())
	default:
		return "global"
	}
}

// IPRateLimiter 基于IP的限流中间件
func IPRateLimiter(fillRate float64, capacity int64, useRedis bool) gin.HandlerFunc {
	return RateLimiter(RateLimiterConfig{
		Type:      IPLimiter,
		FillRate:  fillRate,
		Capacity:  capacity,
		KeyPrefix: "rate_limit:ip",
		Window:    60,
		UseRedis:  useRedis,
	})
}

// UserRateLimiter 基于用户的限流中间件
func UserRateLimiter(fillRate float64, capacity int64, useRedis bool) gin.HandlerFunc {
	return RateLimiter(RateLimiterConfig{
		Type:      UserLimiter,
		FillRate:  fillRate,
		Capacity:  capacity,
		KeyPrefix: "rate_limit:user",
		Window:    60,
		UseRedis:  useRedis,
	})
}

// APIRateLimiter 基于API的限流中间件
func APIRateLimiter(fillRate float64, capacity int64, useRedis bool) gin.HandlerFunc {
	return RateLimiter(RateLimiterConfig{
		Type:      APILimiter,
		FillRate:  fillRate,
		Capacity:  capacity,
		KeyPrefix: "rate_limit:api",
		Window:    60,
		UseRedis:  useRedis,
	})
}

// GlobalRateLimiter 全局限流中间件
func GlobalRateLimiter(fillRate float64, capacity int64, useRedis bool) gin.HandlerFunc {
	return RateLimiter(RateLimiterConfig{
		Type:      GlobalLimiter,
		FillRate:  fillRate,
		Capacity:  capacity,
		KeyPrefix: "rate_limit:global",
		Window:    60,
		UseRedis:  useRedis,
	})
}
